/**
 * 
 */
package cs342.project1.mathNotebook.tests.PKI;


import java.io.FileNotFoundException;
import java.io.IOException;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Date;
import java.util.Random;

import junit.framework.Assert;

import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

import cs342.project1.mathNotebook.PKI.SimpleRSAAlgorithmException;
import cs342.project1.mathNotebook.PKI.SimpleRSAKeyPair;
import cs342.project1.mathNotebook.PKI.SimpleRSAPrivateKey;
import cs342.project1.mathNotebook.PKI.SimpleRSAPublicKey;
import cs342.project1.mathNotebook.db.DBNotebook;
import cs342.project1.mathNotebook.db.DBPerson;
import cs342.project1.mathNotebook.db.DatabaseContext;
import cs342.project1.mathNotebook.db.DatabaseException;
import cs342.project1.mathNotebook.db.DatabaseUnconnectableException;

/**
 * @author aaron
 *
 */
public class SimpleRSATest {


	/**
	 * @throws java.lang.Exception
	 */
	@BeforeClass
	public static void setUpBeforeClass() throws Exception {
	}

	/**
	 * @throws java.lang.Exception
	 */
	@AfterClass
	public static void tearDownAfterClass() throws Exception {
	}

	/**
	 * @throws java.lang.Exception
	 */
	@Before
	public void setUp() throws Exception {
	}

	/**
	 * @throws java.lang.Exception
	 */
	@After
	public void tearDown() throws Exception {
	}
	
	@Test
	public void testKeyPairGenerator() {
		SimpleRSAKeyPair kp = SimpleRSAKeyPair.generateKeyPair();
		Assert.assertTrue("Keys should match", SimpleRSAKeyPair.checkKeyPair(kp.getPublic(), kp.getPrivate()));
	}
		
	@Test
	public void testEncryptDecrypt() {
		Random rng = new Random(new Date().getTime());
		SimpleRSAKeyPair kp = SimpleRSAKeyPair.generateKeyPair();
		for(int i = 1; i <= 245; i++) {
			byte[] input = new byte[i];
			rng.nextBytes(input);
			byte[] output = kp.getPublic().encrypt(input);
			output = kp.getPrivate().decrypt(output);
			
			Assert.assertTrue("Output of decrypt should equal input to encrypt", Arrays.equals(input, output));
		}
	}
	
	@Test
	public void testKeysFromBytes() {
		Random rng = new Random(new Date().getTime());
		for(int i = 0; i < 5; i++) {
			SimpleRSAKeyPair kp = SimpleRSAKeyPair.generateKeyPair();
			byte[] publicBytes = kp.getPublic().getEncoded();
			byte[] privateBytes = kp.getPrivate().getEncoded();
		
			SimpleRSAPublicKey pub = new SimpleRSAPublicKey(publicBytes);
			SimpleRSAPrivateKey pri = new SimpleRSAPrivateKey(privateBytes);
			
			Assert.assertTrue("Keys must pair", SimpleRSAKeyPair.checkKeyPair(pub, pri));
			
			byte[] input = new byte[200];
			rng.nextBytes(input);
			
			byte[] output = pub.encrypt(input);
			output = pri.decrypt(output);
			Assert.assertTrue("Input should equal output", Arrays.equals(input, output));
		}
	}
	
	@Test
	public void testKeysToAndFromFile() throws IOException {
		Random rng = new Random(new Date().getTime());
		SimpleRSAKeyPair kp = SimpleRSAKeyPair.generateKeyPair();
		String filename = "test.pri";
		kp.getPrivate().writeToFile(filename);
		byte[] b = kp.getPublic().getEncoded();
		SimpleRSAPublicKey pub = new SimpleRSAPublicKey(b);
		SimpleRSAPrivateKey pri = new SimpleRSAPrivateKey(filename);
		
		Assert.assertTrue("Keys must pair", SimpleRSAKeyPair.checkKeyPair(pub, pri));
		
		byte[] input = new byte[200];
		rng.nextBytes(input);
		
		byte[] output = pub.encrypt(input);
		output = pri.decrypt(output);
		Assert.assertTrue("Input should equal output", Arrays.equals(input, output));
	}
	
	@Test
	public void sampleUsageTest() throws FileNotFoundException, IOException, DatabaseException, NoSuchAlgorithmException {
		//create a database context
		DatabaseContext context = new DatabaseContext("resources/DatabaseContext-Test.properties");
		
		//CREATE A NEW PERSON
		//first check that the desired username is available
		Assert.assertTrue("username should be available--there are no users", DBPerson.isUsernameAvailable(context, "aaron"));
		//generate random key pair
		SimpleRSAKeyPair keyPair = SimpleRSAKeyPair.generateKeyPair();
		//save private key to file
		keyPair.getPrivate().writeToFile("test.pri");
		//create new person in DB
		DBPerson person = DBPerson.createPerson(context, "aaron", "Aaron Heise", "aaron.heise@gmail.com", keyPair.getPublic().getEncoded());
		//For the sake of this test, invalidate the person so the object will be re-read from the database on next use.
		person.invalidate();
		//END CREATE A NEW PERSON
		
		//CREATE A NEW NOTEBOOK		
		//pretend new notebook
		byte[] createNotebookBytes = new byte[]{'a', 'b', 'c','d','e','f','g'};
		//pretend compress
		byte[] createCompressedNotebookBytes = createNotebookBytes;
		//create a secure random DES key--use SHA1PRNG
		//this would be passed to the DES encryptor
		byte[] desKey = new byte[32]; //256-bit
		SecureRandom secureRandom = SecureRandom.getInstance("SHA1PRNG");
		secureRandom.setSeed(secureRandom.generateSeed(20));
		secureRandom.nextBytes(desKey);
		//pretend DES encrypt
		byte[] createEncryptedCompressedNotebookBytes = createCompressedNotebookBytes;
		//create notebook in database
		DBNotebook notebook = DBNotebook.createNotebook(context, person, new Date(), createEncryptedCompressedNotebookBytes);
		//save the ID of this notebook for later
		int notebookId = notebook.getId();
		//add owner as reader--this is NOT done automatically.
		//first fetch the owner object if you do not already have it.
		DBPerson owner = DBPerson.getPersonByUsername(context, "aaron");
		//deserialize user's public key
		SimpleRSAPublicKey publicKey = new SimpleRSAPublicKey(owner.getPublicKey());
		//encrypt DES key for this user
		byte[] ownerEncryptedDESKey = publicKey.encrypt(desKey);
		//add the key to the readers table
		notebook.addReader(owner, ownerEncryptedDESKey);
		//also add owner as writer
		notebook.addWriter(owner);
		//for future tests, invalidate database object
		owner.invalidate();
		notebook.invalidate();
		//END CREATE A NEW NOTEBOOK
		
		//READ A NOTEBOOK
		//get the reader person
		DBPerson reader = DBPerson.getPersonByUsername(context, "aaron");
		//fetch notebook--pretend that the ID was looked up already.
		//probably the notebook would end up coming from a list created by
		//DBNotebook.getAllNotebooks(context), and either way is fine.
		DBNotebook notebookToRead = DBNotebook.getNotebookById(context, notebookId);
		//get the encrypted DES key for this reader--it will be null if the person is not a reader.
		byte[] readerEncryptedDESKey = notebookToRead.getKeyForReader(reader);
		Assert.assertNotNull("Reader's key should be found", readerEncryptedDESKey);
		//get the reader's private key
		SimpleRSAPrivateKey privateKey = new SimpleRSAPrivateKey("test.pri");
		//decrypt the DES key
		byte[] clearDESKey = privateKey.decrypt(readerEncryptedDESKey);
		Assert.assertTrue("clear DES key should match original DES key", Arrays.equals(desKey, clearDESKey));
		//get the notebook body
		byte[] readerEncryptedCompressedNotebookBytes = notebook.getData();
		//pretend decrypt using clearDESKey, uncompress
		byte[] readerNotebookBytes = readerEncryptedCompressedNotebookBytes;
		Assert.assertTrue("should equal original notebook", Arrays.equals(readerNotebookBytes, createNotebookBytes));
		//for future tests, invalidate database objects
		reader.invalidate();
		notebookToRead.invalidate();
		//END READ A NOTEBOOK
		
		//WRITE A NOTEBOOK
		//(if a writer changes a notebook)
		//get the person
		DBPerson writer = DBPerson.getPersonByUsername(context, "aaron");
		//get the notebook object
		DBNotebook notebookToWrite = DBNotebook.getNotebookById(context, notebookId);
		//check that person has write permission (you would use an if, but since this is a test case...)
		Assert.assertTrue("person should have write access (is owner)", notebookToWrite.hasWriter(writer));
		//a base assumption is that a writer is also a reader, and in this case, we already have
		//the clear DES key in clearDESKey.
		//Use the clear DES key to encrypt the new notebook body (pretending here)
		byte[] writerEncryptedCompressedNotebookBytes = readerNotebookBytes;
		//write the data back out to the database
		notebookToWrite.setData(writerEncryptedCompressedNotebookBytes);
		//since the DES key didn't change, new keys do not need to be disseminated to readers (A Good Thing).
		//invalidate objects for testing purposes
		writer.invalidate();
		notebookToWrite.invalidate();
		//END WRITE A NOTEBOOK
		
		//AUTHENTICATE A USER
		//When a user logs in, they will supply their username and private key.  The software
		//must verify this key pairs with the key in the database for that username.  If it does,
		//the user has successfully authenticated.
		//get the person
		DBPerson user = DBPerson.getPersonByUsername(context, "aaron");
		//get the private key
		SimpleRSAPrivateKey userPrivateKey = new SimpleRSAPrivateKey("test.pri");
		//get the public key
		SimpleRSAPublicKey userPublicKey = new SimpleRSAPublicKey(user.getPublicKey());
		//check (asserting here since it is a test case)
		Assert.assertTrue("private key should match public key", SimpleRSAKeyPair.checkKeyPair(userPublicKey, userPrivateKey));
		//invalidate objects for testing purposes
		user.invalidate();
		//END AUTHENTICATE A USER
		
		
		//cleanup
		DBPerson.clearContext(context);
		DBNotebook.clearContext(context);
		
	}
	

}
